#include "S5ActionUdp.h"
#include <KPackage.h>
#include <KUdpSocket.h>
#include "S5ActionClient.h"
#include "S5Debug.h"
#include "S5Def.h"
#include "S5Server.h"

S5ActionUdp::S5ActionUdp(S5ActionClient *client)
    : m_unpack(false), m_addr(0), m_client(client), m_peer(NULL)
{
    ;
}

void S5ActionUdp::associate(S5ActionUdp *peer)
{
    this->m_peer = peer;
    peer->m_peer = this;
}

void S5ActionUdp::relay(const void *data, size_t length, const KIAddress *dest)
{
    if (dest) {
        server()->sendTo(socket(), *dest, data, length);
    } else {
        server()->sendTo(socket(), m_addr, data, length);
    }
}

void S5ActionUdp::setSourceAddress(const KIAddress &address)
{
    m_unpack = true;
    m_addr.reset(address);
}

bool S5ActionUdp::process(void *data, size_t length, const KIpAddress &src)
{
    if (S5ActionUdp *peer = m_peer) {
        if (m_unpack) {
            if (m_addr.netIp() != 0 && m_addr.netIp() != src.netIp()) return false;
            if (m_addr.netPort() != 0 && m_addr.netPort() != src.netPort()) return false;
            if (length > sizeof(UdpHeaderPrefix)) {
                UdpHeaderPrefix *prefix = reinterpret_cast<UdpHeaderPrefix *>(data);
                if (prefix->rsv == S5_RSV && prefix->frag == S5_FRAG_STAND_ALONE) {
                    switch (prefix->atyp) {
                    case S5_ATYP_IPV4:
                        {
                            if (length > sizeof(UdpHeader)) {
                                UdpHeader *header = reinterpret_cast<UdpHeader *>(data);
                                KIpAddress dest(header->dst_port, header->dst_addr);
                                S5DebugHex(header + 1, length - sizeof(UdpHeader), "Udp send to server", &dest);
                                peer->relay(header + 1, length - sizeof(UdpHeader), &dest);
                                return true;
                            }
                        } break;
                    case S5_ATYP_DOMAINNAME:
                        {
                            uint8_t *byte = reinterpret_cast<uint8_t *>(data) + sizeof(UdpHeaderPrefix);
                            size_t len = *byte++;
                            size_t size = sizeof(UdpHeaderPrefix) + 1 + len + 2;
                            if (length > size) {
                                std::string domain(byte, byte + len);
                                KTcpHostAddress addr(domain.c_str(), NULL);
                                if (!addr.isNull()) {
                                    byte += len;
                                    KIpAddress dest(u_short(byte[0] + (byte[1] << 8)), KIpAddress(addr).netIp());
                                    S5DebugHex(byte + 2, length - size, "Udp send to server", &dest);
                                    peer->relay(byte + 2, length - size, &dest);
                                    return true;
                                }
                            }
                        } break;
                    }
                }
            }
        } else {
            UdpHeader header;
            header.dst_addr = src.netIp();
            header.dst_port = src.netPort();
            KPackage pack(sizeof header + length);
            pack.append(header).append(data, length);
            S5DebugHex(pack.data(), pack.length(), "Udp receive from server", &src);
            peer->relay(pack.data(), pack.length(), NULL);
            return true;
        }
        return false;
    } else {
        return true;
    }
}

void S5ActionUdp::abortImp()
{
    // Close proxy side
    server()->ensureAlive(socket(), false);
    socket()->shutdown();
    socket()->close();
    m_peer = NULL;
    // Close client side
    m_client->abortOnce();
}

void S5ActionUdp::onError(KDef::SockError type, DWORD error)
{
    abortOnce();
}

void S5ActionUdp::onReceived(KKey *&key, void *data, size_t length, bool more)
{
    S5Action::onReceived(key, data, length, more);
    if (more) {
        KDebugMsg("Warning: packet loss");
    } else {
        KUdpSocket *sock = static_cast<KUdpSocket *>(socket());
        if (process(data, length, sock->fromAddress())) {
            return;
        }
    }
    server()->receive(socket());
}

void S5ActionUdp::onSent(KKey *&key, void *data, size_t length)
{
    S5Action::onSent(key, data, length);
    if (S5ActionUdp *peer = m_peer) {
        peer->onPeerRelayed();
    }
}